use std::sync::atomic::AtomicUsize;
use std::sync::mpmc::{Sender, Receiver};
use std::sync::{Arc, Barrier};
use std::thread::JoinHandle;
use std::collections::BTreeMap;

use schedule::error::ScheduleError;
use schedule::exec::*;

use super::mpmc_attr::*;

pub struct MpmcTaskScheduler<T:MpmcTaskAttr> {
    handles:Vec<JoinHandle<()>>,
    task_counter:Arc<AtomicUsize>,
    sched_contorl_receiver:Receiver<()>,
    exec_terminal_sender:Sender<()>,
    dyn_sched_sender:Sender<ProcessControlTask<T>>,
    dyn_exec_receiver:Receiver<ProcessControlTask<T>>,
    specific_sched_map:BTreeMap<String, usize>,
    specific_sched_pair_vec:Vec<(Sender<ProcessControlTask<T>>, Receiver<ProcessControlTask<T>>)>,
    dyn_woker_num:usize,
    init_task:Option<T>,
    error_collector:ErrorCollector<ScheduleError>,
}

impl<T:MpmcTaskAttr> MpmcTaskScheduler<T> {
    pub fn new(init_task:T) -> Self {
        let error_collector = ErrorCollector::new();
        let  (dyn_sched_sender, dyn_exec_receiver) = std::sync::mpmc::channel::<ProcessControlTask<T>>();
        let (exec_terminal_sender, sched_contorl_receiver) = std::sync::mpmc::channel::<()>();
        let mut sched = Self { 
            handles:Vec::new(), 
            task_counter:Arc::new(AtomicUsize::new(0)),
            sched_contorl_receiver, exec_terminal_sender,
            error_collector, dyn_sched_sender, dyn_exec_receiver,
            dyn_woker_num:0,
            init_task:Some(init_task),
            specific_sched_map:BTreeMap::new(),
            specific_sched_pair_vec:Vec::new(),
        };
        T::predefine_sender(&mut sched);
        sched
    }

    pub fn task_counter(&self) -> Arc<AtomicUsize> {
        self.task_counter.clone()
    }

    pub fn register_specific_executor_name<Q:ToString>(&mut self, name:Q) {
        let dyn_msg_pair = std::sync::mpmc::channel::<ProcessControlTask<T>>();
        self.specific_sched_map.insert(name.to_string(), self.specific_sched_pair_vec.len());
        self.specific_sched_pair_vec.push(dyn_msg_pair);
    }

    pub fn get_exec_terminal_sender(&self) -> Sender<()> {
        self.exec_terminal_sender.clone()
    }

    pub fn get_specific_sched_sender<Q:ToString>(&self, name:Q) -> Option<Sender<ProcessControlTask<T>>> {
        self.specific_sched_map.get(&name.to_string()).map(|&index|self.specific_sched_pair_vec[index].0.clone())
    }

    pub fn get_dyn_sched_sender(&self) -> Sender<ProcessControlTask<T>> {
        self.dyn_sched_sender.clone()
    }

    pub fn join(&mut self) {
        while !self.handles.is_empty() {
            self.handles.pop().unwrap().join().expect("join handle error");
        }
    }
}

impl<T:MpmcTaskAttr> TaskScheduler<T> for MpmcTaskScheduler<T> {
    type Sender = T::SenderGroup;
    type Receiver = Receiver<ProcessControlTask<T>>;
    type Error = ScheduleError;
    fn send_with_count(sender:&Self::Sender, task:T, task_counter:&AtomicUsize) {
        task_counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        T::group_send(sender, ProcessControlTask::Task(task));
    }

    fn recv(recv:&Self::Receiver) -> Option<T> {
        let process_contorl_task = recv.recv().expect("mpmc receive");
        match process_contorl_task {
            ProcessControlTask::Task(task) => { return Some(task); },
            ProcessControlTask::Break(barrier) => { barrier.wait(); return None; },
        }
    }

    fn register_dyn_executor<TE:TaskExecutor<T>>(&mut self, executor:TE) {
        let error_collector = self.error_collector.clone();
        let sender_group = T::get_sender_group_from(self).expect("sender missed");
        let exec_recv = self.dyn_exec_receiver.clone();
        let task_counter = self.task_counter.clone();
        self.dyn_woker_num += 1;

        let handle = std::thread::spawn(move || {
            let run_result = executor.run::<Self>(exec_recv, sender_group, task_counter.clone());
            error_collector.collect(run_result.map_err(|e|format!("{e:?}").into()))
        });
        self.handles.push(handle);
    }

    fn register_specific_executor<TE:TaskExecutor<T>>(&mut self, name:String, executor:TE) -> bool {
        let msg_pair_index = self.specific_sched_map.get(&name);
        if msg_pair_index.is_none() { return false; }
        let msg_pair = &self.specific_sched_pair_vec[*msg_pair_index.unwrap()];
        let exec_recv = msg_pair.1.clone();
        let error_collector = self.error_collector.clone();
        
        let sender_group = T::get_sender_group_from(self).expect("sender missed");
        let task_counter = self.task_counter.clone();
        let handle = std::thread::spawn(move || {
            let run_result = executor.run::<Self>(exec_recv, sender_group, task_counter.clone());
            error_collector.collect(run_result.map_err(|e|format!("{e:?}").into()))
        });
        self.handles.push(handle);
        true
    }

    fn schedule(&mut self) -> Result<(), Self::Error> {
        let init_task = self.init_task.take().unwrap();
        self.task_counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        T::send_init_task(self, init_task);
        let _ = self.sched_contorl_receiver.recv().expect("sched contorl error");
        let worker_num = self.handles.len();
        let barrier = Arc::new(Barrier::new(worker_num));
        for (specific_exec_sender, _) in self.specific_sched_pair_vec.iter() {
            specific_exec_sender.send(ProcessControlTask::Break(barrier.clone())).expect("terminale barrier send error");
        }
        for _ in 0..self.dyn_woker_num {
            self.dyn_sched_sender.send(ProcessControlTask::Break(barrier.clone())).expect("terminale barrier send error");
        }
        self.join();
        Ok(())
    }

    fn terminal(sender:&Self::Sender) {
        let exec_terminal_sender = T::get_exec_terminal_sender(sender);
        exec_terminal_sender.send(()).expect("terminale error");
    }

    fn error_collector(&self) -> ErrorCollector<Self::Error> {
        self.error_collector.clone()
    }
}